//
//  H264StreamDecoder.m
//  testdemo
//
//  Created by jiang on 2017/6/21.
//  Copyright © 2017年 jiang. All rights reserved.
//

#import "H264HardwareStreamDecoder.h"
#import <VideoToolbox/VideoToolbox.h>

//暂时为4bytes
static const uint8_t StartCode[4] = {0x00, 0x00, 0x00, 0x01};

@interface H264HardwareStreamDecoder (){
    VTDecompressionSessionRef decompressionSession;
    CMVideoFormatDescriptionRef _decoderFormatDescription;
    
    uint32_t _spsSize;
    uint32_t _ppsSize;
    
    uint8_t *_sps;
    uint8_t *_pps;
    
    //输入 数据 及 长度，外部输入
    uint8_t *_inputBuffer;
    NSInteger _inputSize;
    
    //解析出完整的 一帧 和 长度
    uint8_t *_packetBuffer;
    NSInteger _packetSize;
    
    //上次remaining data
    uint8_t *_remainingBuffer;
    NSInteger _remainingSize;
}

@property (nonatomic) BOOL isBackGround;

@end

@implementation H264HardwareStreamDecoder

#pragma mark - life cycle
- (instancetype)init {
    if (self = [super init]) {
        
        [[NSNotificationCenter defaultCenter] addObserver:self selector:@selector(willEnterBackground:) name:UIApplicationWillResignActiveNotification object:nil];
        [[NSNotificationCenter defaultCenter] addObserver:self selector:@selector(willEnterForeground:) name:UIApplicationDidBecomeActiveNotification object:nil];
    }
    return self;
}

- (void)dealloc {
    
    if (_remainingBuffer) {
        free(_remainingBuffer);
        _remainingBuffer = NULL;
    }
    
    if (_packetBuffer) {
        free(_packetBuffer);
        _packetBuffer = NULL;
    }
    
    if (_sps) {
        free(_sps);
        _sps = NULL;
    }
    _spsSize = 0;
    
    if (_pps) {
        free(_pps);
        _pps = NULL;
    }
    _ppsSize = 0;
    
    if(_decoderFormatDescription) {
        CFRelease(_decoderFormatDescription);
        _decoderFormatDescription = NULL;
    }
    
    if (decompressionSession) {
        VTDecompressionSessionInvalidate(decompressionSession);
        CFRelease(decompressionSession);
        decompressionSession = NULL;
    }
    [[NSNotificationCenter defaultCenter] removeObserver:self];
}

#pragma mark - notification
- (void)willEnterBackground:(NSNotification *)notification {
    _isBackGround = YES;
}

- (void)willEnterForeground:(NSNotification *)notification {
    _isBackGround = NO;
}

#pragma mark - processing network or local stream

/**
 处理接收到的流数据，不一定是完整的NALU，此处处理完一帧直接break，后续先解压一帧
 */
- (void)_processInputBufferStream {
    if (_inputBuffer == NULL) {
        return;
    }
    
    //inputbuffer start with 00 00 00 01
    if (memcmp(_inputBuffer, StartCode, 4) == 0) {
        if (_inputSize > 4) {
            uint8_t *pStart = _inputBuffer + 4;
            uint8_t *pEnd = _inputBuffer + _inputSize;
            
            while (pStart != pEnd) {
                if (memcmp(pStart - 3, StartCode, 4) == 0) {
                    _packetSize = pStart - _inputBuffer - 3;
                    _packetBuffer = realloc(_packetBuffer, _packetSize);
                    if (!_packetBuffer) { // memory alloc failed
                        return;
                    }
                    
                    memcpy(_packetBuffer, _inputBuffer, _packetSize);
                    //move inputbuffer
                    memmove(_inputBuffer, pStart + 1, pEnd - (pStart + 1));
                    _inputSize = pEnd - (pStart + 1);
                    
                    //处理完整一帧
                    [self _startDecodingNALUPacket];
                } else {
                    pStart++;
                }
            }
            
            /**
             *  inputBuffer剩下没处理完的（而且不足以组装 完整一帧），copy到_remainingBuffer，等待下次新的inputBuffer输入，
             *  再次遇到StartCode，copy新的到 拼接_remainingBuffer(range：0--StartCode Index)
             **/
            if (_inputSize > 0) {
                _remainingSize = _inputSize;
                _remainingBuffer = realloc(_remainingBuffer, _remainingSize);
                if (!_remainingBuffer) { // memory alloc failed
                    return;
                }
                memcpy(_remainingBuffer, _inputBuffer, _remainingSize);
            }
        }
    } else {
        //不以StartCode 开头
        if (!_remainingBuffer) {
            NSLog(@"_remainingBuffer 暂时不能为空");
            return;
        }
        uint8_t *pStart = _inputBuffer;
        uint8_t *pEnd = _inputBuffer + _inputSize;
        
        while (pStart != pEnd) {
            if (memcmp(pStart - 3, StartCode, 4) == 0) {
                // 拼接上次数据
                _packetSize = pStart - _inputBuffer - 3 + _remainingSize;
                _packetBuffer = realloc(_packetBuffer, _packetSize);
                if (!_packetBuffer) { // memory alloc failed
                    return;
                }
                
                memcpy(_packetBuffer, _remainingBuffer, _remainingSize);
                memcpy(_packetBuffer + _remainingSize, _inputBuffer, (_packetSize - _remainingSize));
                //move inputbuffer
                memmove(_inputBuffer, pStart + 1, pEnd - (pStart + 1));
                _inputSize = pEnd - (pStart + 1);
                
                //处理完整一帧
                [self _startDecodingNALUPacket];
            } else {
                pStart++;
            }
        }
        
        if (_inputSize > 0) {
            _remainingSize = _inputSize;
            _remainingBuffer = realloc(_remainingBuffer, _remainingSize);
            if (!_remainingBuffer) { // memory alloc failed
                return;
            }
            memcpy(_remainingBuffer, _inputBuffer, _remainingSize);
        }
    }
}

#pragma mark - decode
- (void)_resetDecompressionSession {
    if (decompressionSession) {
        VTDecompressionSessionInvalidate(decompressionSession);
        CFRelease(decompressionSession);
        decompressionSession = NULL;
    }
    
    const uint8_t* const parameterSetPointers[2] = { _sps, _pps };
    const size_t parameterSetSizes[2] = { _spsSize, _ppsSize };
    OSStatus status = CMVideoFormatDescriptionCreateFromH264ParameterSets(kCFAllocatorDefault,
                                                                          2, //param count
                                                                          parameterSetPointers,
                                                                          parameterSetSizes,
                                                                          4, //nal start code size
                                                                          &_decoderFormatDescription);
    
    if(status == noErr) {
        CFDictionaryRef attrs = NULL;
        const void *keys[] = { kCVPixelBufferPixelFormatTypeKey };
        //      kCVPixelFormatType_420YpCbCr8Planar is YUV420
        //      kCVPixelFormatType_420YpCbCr8BiPlanarFullRange is NV12
        uint32_t v = kCVPixelFormatType_420YpCbCr8BiPlanarFullRange;
        const void *values[] = { CFNumberCreate(NULL, kCFNumberSInt32Type, &v) };
        attrs = CFDictionaryCreate(NULL, keys, values, 1, NULL, NULL);
        
        VTDecompressionOutputCallbackRecord callBackRecord;
        callBackRecord.decompressionOutputCallback = decompressionSessionDecodeFrameCallback;
        callBackRecord.decompressionOutputRefCon = (__bridge void *)self;
        
        status = VTDecompressionSessionCreate(kCFAllocatorDefault,
                                              _decoderFormatDescription,
                                              NULL,
                                              attrs,
                                              &callBackRecord,
                                              &decompressionSession);
        CFRelease(attrs);
    } else {
        NSLog(@"reset decoder session failed status=%d", (int)status);
    }
}

/**
 decoding NALU buffer
 */
- (void)_decode {
    if (decompressionSession) {
        // compress bitstream to CMBlockBufferRef
        CMBlockBufferRef blockBuffer = NULL;
        OSStatus status = CMBlockBufferCreateWithMemoryBlock(kCFAllocatorDefault,
                                                             (void *)_packetBuffer,
                                                             _packetSize,
                                                             kCFAllocatorNull,
                                                             NULL,
                                                             0,
                                                             _packetSize,
                                                             0,
                                                             &blockBuffer);
        
        if (status == kCMBlockBufferNoErr) {
            //CMBlockBufferRef to CMSampleBufferRef
            CMSampleBufferRef sampleBuffer = NULL;
            const size_t sampleSizeArray[] = { _packetSize };
            
            status = CMSampleBufferCreate(kCFAllocatorDefault,
                                          blockBuffer,
                                          true,
                                          NULL,
                                          NULL,
                                          _decoderFormatDescription,
                                          1, 0, NULL, 1,
                                          sampleSizeArray,
                                          &sampleBuffer);
            
            //            status = CMSampleBufferCreateReady(kCFAllocatorDefault,
            //                                               blockBuffer,
            //                                               _decoderFormatDescription,
            //                                               1, 0, NULL, 1,
            //                                               sampleSizeArray,
            //                                               &sampleBuffer);
            
            if (status == kCMBlockBufferNoErr && sampleBuffer) {
                
                VTDecodeFrameFlags flags = 0;
                VTDecodeInfoFlags flagOut;
                OSStatus decodeStatus = VTDecompressionSessionDecodeFrame(decompressionSession,
                                                                          sampleBuffer,
                                                                          flags,
                                                                          &sampleBuffer,
                                                                          &flagOut);
                
                if(decodeStatus == kVTInvalidSessionErr) {
                    NSLog(@"Invalid session, reset decoder session");
                } else if(decodeStatus == kVTVideoDecoderBadDataErr) {
                    NSLog(@"decode failed status=%d(Bad data)", decodeStatus);
                } else if(decodeStatus != noErr) {
                    NSLog(@"decode failed status=%d", decodeStatus);
                }
                
                CFRelease(sampleBuffer);
            }
        }
        
        if (blockBuffer) {
            CFRelease(blockBuffer);
        }
    }
}

/**
 处理完原始buffer后，该函数 开始处理完整的NALU帧
 */
- (void)_startDecodingNALUPacket {
    if (_packetBuffer == NULL || _packetSize <= 0) {
        return;
    }
    
    //replace start code with size
    uint32_t naluSize = (uint32_t)(_packetSize - 4);
    uint32_t *pNaluSize = (uint32_t *)_packetBuffer;
    *pNaluSize = CFSwapInt32HostToBig(naluSize);
    
    int naluType = _packetBuffer[4] & 0x1F;
    switch (naluType) {
        case 0x07:
        {
            // sps
            _spsSize = (uint32_t)(_packetSize - 4);
            if (_sps) {
                free(_sps);
                _sps = NULL;
            }
            _sps = malloc(_spsSize);
            memcpy(_sps, pNaluSize + 4, _spsSize);
        }
            break;
        case 0x08:
        {
            // pps
            _ppsSize = (uint32_t)(_packetSize - 4);
            if (_pps) {
                free(_pps);
                _pps = NULL;
            }
            _pps = malloc(_ppsSize);
            memcpy(_pps, pNaluSize + 4, _ppsSize);
        }
            break;
        case 0x05:
        {
            // IDR / I frame
            [self _resetDecompressionSession];
            [self _decode];
        }
            break;
        case 0x01:
        {
            // P frame
            [self _decode];
        }
            break;
        default:
        {
            // B Frame or other
            [self _decode];
        }
            break;
    }
}

- (void)endDecoding {
    if (decompressionSession) {
        VTDecompressionSessionInvalidate(decompressionSession);
        CFRelease(decompressionSession);
        decompressionSession = NULL;
    }
}

- (void)decode:(uint8_t *)inputBuffer inputSize:(NSInteger)inputSize {
    if (_isBackGround) {
        return;
    }
    _inputBuffer = inputBuffer;
    _inputSize = inputSize;
    
    [self _processInputBufferStream];
}

#pragma mark - decode callback
static void decompressionSessionDecodeFrameCallback(void * decompressionOutputRefCon,
                                                    void * sourceFrameRefCon,
                                                    OSStatus status,
                                                    VTDecodeInfoFlags infoFlags,
                                                    CVImageBufferRef imageBuffer,
                                                    CMTime presentationTimeStamp,
                                                    CMTime presentationDuration ) {
    NSLog(@"%s, status is %@", __FUNCTION__, (status == noErr) ? @"ok " : @"error");
    CVPixelBufferLockBaseAddress(imageBuffer, 0);
    
    //commit imageBuffer to OpenGL ES
    
    CVPixelBufferUnlockBaseAddress(imageBuffer, 0);
}

@end
